import json
import requests
from sseclient import SSEClient


class MCPSSEClient:
    def __init__(self, server_url: str):
        self.server_url = server_url

    def sync_call(self, params: dict) -> dict:
        """
        同步调用MCP服务端，等待完全完成后再返回

        :param params: 调用参数
        :return: 最终结果
        """
        # 准备请求头
        headers = {'Accept': 'text/event-stream'}
        event_url = f"{self.server_url}?params={requests.utils.quote(json.dumps(params))}"

        try:
            # 创建SSE客户端
            response = requests.get(event_url, stream=True, headers=headers)
            client = SSEClient(response)

            final_result = None

            # 正确迭代SSE事件
            for event in client.events():
                if event.event == 'error':
                    error_data = json.loads(event.data)
                    raise Exception(f"Server error: {error_data.get('message')}")

                if event.event == 'progress':
                    progress_data = json.loads(event.data)
                    print(f"Progress: {progress_data.get('progress', 0)}%")

                if event.event == 'complete':
                    final_result = json.loads(event.data)
                    break

            if final_result is None:
                raise Exception("Server did not send completion event")

            return final_result

        except requests.exceptions.RequestException as e:
            raise Exception(f"Connection failed: {str(e)}")
        except json.JSONDecodeError as e:
            raise Exception(f"Invalid server response: {str(e)}")


# 使用示例
if __name__ == "__main__":
    client = MCPSSEClient("http://localhost:8081/sse")
    topic = "物流管理系统"
    content = "软件简介   物流信息管理系统主要用于物流企业对运输任务、车辆、司机及货物等信息的全面管理，实现物流过程的信息化和智能化。"
    args = {"topic": topic, "content": content}
    try:
        result = client.sync_call({
            "task": "plantuml_generator",
            "topic": topic, "content": content
        })
        print("任务完成，结果:", result)
        # 在这里添加后续处理逻辑

    except Exception as e:
        print("调用失败:", str(e))